/*
 * @file arma_reduce.h
 *
 *  
 *
 * @author Dahua Lin
 */

#ifndef _MSC_VER
#pragma once
#endif

#ifndef GMLPP_ARMA_REDUCE_H_
#define GMLPP_ARMA_REDUCE_H_

#include "arma_base.h"

namespace gml
{
	/********************************************
	 *
	 *  map_accu / map_sum
	 *
	 ********************************************/

	template<class Fun, typename T>
	inline
	typename Fun::result_type map_accu(Fun fun, const Mat<T>& X)
	{
		uword n = X.n_elem;
		const T *x = X.memptr();

		typename Fun::result_type s;
		for (uword i = 0; i < n; ++i) s += fun(x[i]);
		return s;
	}

	template<class Fun, typename T1, typename T2>
	inline
	typename Fun::result_type map_accu(Fun fun, const Mat<T1>& X1, const Mat<T2>& X2)
	{
		check_arg( is_same_size(X1, X2) );

		uword n = X1.n_elem;
		const T1 *x1 = X1.memptr();
		const T1 *x2 = X2.memptr();

		typename Fun::result_type s;
		for (uword i = 0; i < n; ++i) s += fun(x1[i], x2[i]);
		return s;
	}

	template<class Fun, typename T>
	inline
	typename Fun::result_type map_sum(Fun fun, const Row<T>& X)
	{
		return map_accu(fun, X);
	}

	template<class Fun, typename T>
	inline
	typename Fun::result_type map_sum(Fun fun, const Col<T>& X)
	{
		return map_accu(fun, X);
	}


	template<class Fun, typename T1, typename T2>
	inline
	typename Fun::result_type map_sum(Fun fun, const Row<T1>& X1, const Row<T2>& X2)
	{
		return map_accu(fun, X1, X2);
	}

	template<class Fun, typename T1, typename T2>
	inline
	typename Fun::result_type map_sum(Fun fun, const Col<T1>& X1, const Col<T2>& X2)
	{
		return map_accu(fun, X1, X2);
	}


	template<class Fun, typename T, typename TR>
	inline
	void map_sum(Fun fun, const Mat<T>& X, uword dim, Mat<TR>& R)
	{
		uword m = X.n_rows;
		uword n = X.n_cols;

		if (dim == 0)  // per-column sum
		{
			R.set_size(1, n);
			TR *r = R.memptr();

			for (uword j = 0; j < n; ++j)
			{
				const T *x = X.colptr(j);

				TR s = 0;
				for (uword i = 0; i < m; ++i) s += fun(x[i]);
				r[j] = s;
			}
		}
		else  // per-row sum
		{
			R.zeros(m, 1);
			TR *r = R.memptr();

			for (uword j = 0; j < n; ++j)
			{
				const T *x = X.colptr(j);
				for (uword i = 0; i < m; ++i) r[i] += fun(x[i]);
			}
		}
	}


	template<class Fun, typename T1, typename T2, typename TR>
	inline
	void map_sum(Fun fun, const Mat<T1>& X1, const Mat<T2>& X2, uword dim, Mat<TR>& R)
	{
		check_arg( is_same_size(X1, X2) );

		uword m = X1.n_rows;
		uword n = X2.n_cols;

		if (dim == 0)  // per-column sum
		{
			R.set_size(1, n);
			TR *r = R.memptr();

			for (uword j = 0; j < n; ++j)
			{
				const T1 *x1 = X1.colptr(j);
				const T2 *x2 = X2.colptr(j);

				TR s = 0;
				for (uword i = 0; i < m; ++i) s += fun(x1[i], x2[i]);
				r[j] = s;
			}
		}
		else  // per-row sum
		{
			R.zeros(m, 1);
			TR *r = R.memptr();

			for (uword j = 0; j < n; ++j)
			{
				const T1 *x1 = X1.colptr(j);
				const T2 *x2 = X2.colptr(j);
				for (uword i = 0; i < m; ++i) r[i] += fun(x1[i], x2[i]);
			}
		}
	}


	template<class Fun, typename T>
	inline
	Mat<typename Fun::result_type> map_sum(Fun fun, const Mat<T>& X, uword dim)
	{
		Mat<typename Fun::result_type> R;
		map_sum(fun, X, dim, R);
		return R;
	}

	template<class Fun, typename T1, typename T2>
	inline
	Mat<typename Fun::result_type> map_sum(Fun fun, const Mat<T1>& X1, const Mat<T2>& X2, uword dim)
	{
		check_arg( is_same_size(X1, X2) );

		Mat<typename Fun::result_type> R;
		map_sum(fun, X1, X2, dim, R);
		return R;
	}

}



#endif /* ARMA_REDUCE_H_ */
